use base64::prelude::*;
use clap::{Parser, Subcommand};
use crypto::{
    aes, blockmodes,
    buffer::{BufferResult, ReadBuffer, RefReadBuffer, RefWriteBuffer, WriteBuffer},
    symmetriccipher::SymmetricCipherError,
};
use rand::{thread_rng, Rng};

/// AES 编解码工具
#[derive(Debug, Parser)]
struct Cli {
    #[command(subcommand)]
    command: Commands,
}

#[derive(Debug, Subcommand)]
enum Commands {
    /// 编码
    Encode { secret: String, content: String },
    /// 解码
    Decode { secret: String, content: String },
}

fn main() {
    let cli = Cli::parse();
    println!("{:?}", cli);

    match cli.command {
        Commands::Encode { secret, content } => match encrypt(secret, content) {
            Ok(res) => println!("{}", res),
            Err(e) => println!("加密失败, err: {:?}", e),
        },
        Commands::Decode { secret, content } => match decrypt(secret, content) {
            Ok(res) => println!("{}", res),
            Err(e) => println!("解密失败, err: {:?}", e),
        },
    }
}

fn gen_iv() -> [u8; 16] {
    let mut rng = thread_rng();
    let mut iv: [u8; 16] = [0; 16];
    rng.fill(&mut iv);
    iv
}

fn encrypt(secret: String, content: String) -> Result<String, SymmetricCipherError> {
    let iv = gen_iv();
    let mut encryptor = aes::cbc_encryptor(
        aes::KeySize::KeySize256,
        secret.as_bytes(),
        &iv,
        blockmodes::PkcsPadding,
    );
    let mut final_result = Vec::from(iv);
    let mut read_buffer = RefReadBuffer::new(content.as_ref());
    let mut buffer = [0; 4096];
    let mut write_buffer = RefWriteBuffer::new(&mut buffer);

    loop {
        let result = encryptor.encrypt(&mut read_buffer, &mut write_buffer, true)?;

        final_result.extend(
            write_buffer
                .take_read_buffer()
                .take_remaining()
                .iter()
                .copied(),
        );

        match result {
            BufferResult::BufferUnderflow => break,
            BufferResult::BufferOverflow => {}
        }
    }

    Ok(BASE64_STANDARD.encode(final_result))
}

fn decrypt(secret: String, content: String) -> Result<String, SymmetricCipherError> {
    let mut iv = BASE64_STANDARD.decode(content).unwrap();
    let data = iv.split_off(16);

    let mut decryptor = aes::cbc_decryptor(
        aes::KeySize::KeySize256,
        secret.as_bytes(),
        &iv,
        blockmodes::PkcsPadding,
    );
    let mut final_result = Vec::<u8>::new();
    let mut read_buffer = RefReadBuffer::new(data.as_ref());
    let mut buffer = [0; 4096];
    let mut write_buffer = RefWriteBuffer::new(&mut buffer);

    loop {
        let result = decryptor.decrypt(&mut read_buffer, &mut write_buffer, true)?;

        final_result.extend(
            write_buffer
                .take_read_buffer()
                .take_remaining()
                .iter()
                .copied(),
        );

        match result {
            BufferResult::BufferUnderflow => break,
            BufferResult::BufferOverflow => {}
        }
    }

    Ok(String::from_utf8(final_result).unwrap())
}

#[test]
fn test_encrypt() {
    let res = encrypt(
        "39dc02e8bfb44e57bae3ea595086701e".into(),
        "1234567890abcdefg".into(),
    );
    println!("{:?}", res);
}

#[test]
fn test_decrypt() {
    let res = decrypt(
        "39dc02e8bfb44e57bae3ea595086701e".into(),
        "P2N3EVHcSyH+JeWFdNakhMDsH1Pt/ZFKnNC0sPpC4VtVjhqu1T5dXClUOVcpTCup".into(),
    );
    println!("{:?}", res);
}
